// This file is part of Substrate.

// Copyright (C) 2020-2021 Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: Apache-2.0

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// 	http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Helpers for offchain worker election.

use crate::{
    Call, CompactAssignments, Config, ElectionSize, Module, NominatorIndex, Nominators, OffchainAccuracy,
    ValidatorIndex, WeightInfo,
};
use codec::Decode;
use frame_support::{traits::Get, weights::Weight, IterableStorageMap};
use frame_system::offchain::SubmitTransaction;
use sp_npos_elections::{
    reduce, to_support_map, Assignment, CompactSolution, ElectionResult, ElectionScore, EvaluateSupport,
    ExtendedBalance,
};
use sp_runtime::{offchain::storage::StorageValueRef, traits::TrailingZeroInput, RuntimeDebug};
use sp_std::{convert::TryInto, prelude::*};

/// Error types related to the offchain election machinery.
#[derive(RuntimeDebug)]
pub enum OffchainElectionError {
    /// election returned None. This means less candidate that minimum number of needed
    /// validators were present. The chain is in trouble and not much that we can do about it.
    ElectionFailed,
    /// Submission to the transaction pool failed.
    PoolSubmissionFailed,
    /// The snapshot data is not available.
    SnapshotUnavailable,
    /// Error from npos-election crate. This usually relates to compact operation.
    InternalElectionError(sp_npos_elections::Error),
    /// One of the computed winners is invalid.
    InvalidWinner,
    /// A nominator is not available in the snapshot.
    NominatorSnapshotCorrupt,
}

impl From<sp_npos_elections::Error> for OffchainElectionError {
    fn from(e: sp_npos_elections::Error) -> Self {
        Self::InternalElectionError(e)
    }
}

/// Storage key used to store the persistent offchain worker status.
pub(crate) const OFFCHAIN_HEAD_DB: &[u8] = b"parity/staking-election/";
/// The repeat threshold of the offchain worker. This means we won't run the offchain worker twice
/// within a window of 5 blocks.
pub(crate) const OFFCHAIN_REPEAT: u32 = 5;
/// Default number of blocks for which the unsigned transaction should stay in the pool
pub(crate) const DEFAULT_LONGEVITY: u64 = 25;

/// Checks if an execution of the offchain worker is permitted at the given block number, or not.
///
/// This essentially makes sure that we don't run on previous blocks in case of a re-org, and we
/// don't run twice within a window of length [`OFFCHAIN_REPEAT`].
///
/// Returns `Ok(())` if offchain worker should happen, `Err(reason)` otherwise.
pub(crate) fn set_check_offchain_execution_status<T: Config>(now: T::BlockNumber) -> Result<(), &'static str> {
    let storage = StorageValueRef::persistent(&OFFCHAIN_HEAD_DB);
    let threshold = T::BlockNumber::from(OFFCHAIN_REPEAT);

    let mutate_stat = storage.mutate::<_, &'static str, _>(|maybe_head: Option<Option<T::BlockNumber>>| {
        match maybe_head {
            Some(Some(head)) if now < head => Err("fork."),
            Some(Some(head)) if now >= head && now <= head + threshold => Err("recently executed."),
            Some(Some(head)) if now > head + threshold => {
                // we can run again now. Write the new head.
                Ok(now)
            },
            _ => {
                // value doesn't exists. Probably this node just booted up. Write, and run
                Ok(now)
            },
        }
    });

    match mutate_stat {
        // all good
        Ok(Ok(_)) => Ok(()),
        // failed to write.
        Ok(Err(_)) => Err("failed to write to offchain db."),
        // fork etc.
        Err(why) => Err(why),
    }
}

/// The internal logic of the offchain worker of this module. This runs the phragmen election,
/// compacts and reduces the solution, computes the score and submits it back to the chain as an
/// unsigned transaction, without any signature.
pub(crate) fn compute_offchain_election<T: Config>() -> Result<(), OffchainElectionError> {
    let iters = get_balancing_iters::<T>();
    // compute raw solution. Note that we use `OffchainAccuracy`.
    let ElectionResult { winners, assignments } =
        <Module<T>>::do_phragmen::<OffchainAccuracy>(iters).ok_or(OffchainElectionError::ElectionFailed)?;

    // process and prepare it for submission.
    let (winners, compact, score, size) =
        prepare_submission::<T>(assignments, winners, true, T::OffchainSolutionWeightLimit::get())?;

    crate::log!(info, "💸 prepared a seq-phragmen solution with {} balancing iterations and score {:?}", iters, score,);

    // defensive-only: current era can never be none except genesis.
    let current_era = <Module<T>>::current_era().unwrap_or_default();

    // send it.
    let call = Call::submit_election_solution_unsigned(winners, compact, score, current_era, size).into();

    SubmitTransaction::<T, Call<T>>::submit_unsigned_transaction(call)
        .map_err(|_| OffchainElectionError::PoolSubmissionFailed)
}

/// Get a random number of iterations to run the balancing.
///
/// Uses the offchain seed to generate a random number.
pub fn get_balancing_iters<T: Config>() -> usize {
    match T::MaxIterations::get() {
        0 => 0,
        max @ _ => {
            let seed = sp_io::offchain::random_seed();
            let random = <u32>::decode(&mut TrailingZeroInput::new(seed.as_ref()))
                .expect("input is padded with zeroes; qed")
                % max.saturating_add(1);
            random as usize
        },
    }
}

/// Find the maximum `len` that a compact can have in order to fit into the block weight.
///
/// This only returns a value between zero and `size.nominators`.
pub fn maximum_compact_len<W: crate::WeightInfo>(winners_len: u32, size: ElectionSize, max_weight: Weight) -> u32 {
    use sp_std::cmp::Ordering;

    if size.nominators < 1 {
        return size.nominators;
    }

    let max_voters = size.nominators.max(1);
    let mut voters = max_voters;

    // helper closures.
    let weight_with = |voters: u32| -> Weight {
        W::submit_solution_better(size.validators.into(), size.nominators.into(), voters, winners_len)
    };

    let next_voters = |current_weight: Weight, voters: u32, step: u32| -> Result<u32, ()> {
        match current_weight.cmp(&max_weight) {
            Ordering::Less => {
                let next_voters = voters.checked_add(step);
                match next_voters {
                    Some(voters) if voters < max_voters => Ok(voters),
                    _ => Err(()),
                }
            },
            Ordering::Greater => voters.checked_sub(step).ok_or(()),
            Ordering::Equal => Ok(voters),
        }
    };

    // First binary-search the right amount of voters
    let mut step = voters / 2;
    let mut current_weight = weight_with(voters);
    while step > 0 {
        match next_voters(current_weight, voters, step) {
            // proceed with the binary search
            Ok(next) if next != voters => {
                voters = next;
            },
            // we are out of bounds, break out of the loop.
            Err(()) => break,
            // we found the right value - early exit the function.
            Ok(next) => return next,
        }
        step = step / 2;
        current_weight = weight_with(voters);
    }

    // Time to finish.
    // We might have reduced less than expected due to rounding error. Increase one last time if we
    // have any room left, the reduce until we are sure we are below limit.
    while voters + 1 <= max_voters && weight_with(voters + 1) < max_weight {
        voters += 1;
    }
    while voters.checked_sub(1).is_some() && weight_with(voters) > max_weight {
        voters -= 1;
    }

    debug_assert!(
        weight_with(voters.min(size.nominators)) <= max_weight,
        "weight_with({}) <= {}",
        voters.min(size.nominators),
        max_weight,
    );
    voters.min(size.nominators)
}

/// Greedily reduce the size of the a solution to fit into the block, w.r.t. weight.
///
/// The weight of the solution is foremost a function of the number of voters (i.e.
/// `compact.len()`). Aside from this, the other components of the weight are invariant. The number
/// of winners shall not be changed (otherwise the solution is invalid) and the `ElectionSize` is
/// merely a representation of the total number of stakers.
///
/// Thus, we reside to stripping away some voters. This means only changing the `compact` struct.
///
/// Note that the solution is already computed, and the winners are elected based on the merit of
/// teh entire stake in the system. Nonetheless, some of the voters will be removed further down the
/// line.
///
/// Indeed, the score must be computed **after** this step. If this step reduces the score too much,
/// then the solution will be discarded.
pub fn trim_to_weight<T: Config, FN>(
    maximum_allowed_voters: u32,
    mut compact: CompactAssignments,
    nominator_index: FN,
) -> Result<CompactAssignments, OffchainElectionError>
where
    for<'r> FN: Fn(&'r T::AccountId) -> Option<NominatorIndex>,
{
    match compact.voter_count().checked_sub(maximum_allowed_voters as usize) {
        Some(to_remove) if to_remove > 0 => {
            // grab all voters and sort them by least stake.
            let balance_of = <Module<T>>::slashable_balance_of_fn();
            let mut voters_sorted =
                <Nominators<T>>::iter().map(|(who, _)| (who.clone(), balance_of(&who))).collect::<Vec<_>>();
            voters_sorted.sort_by_key(|(_, y)| *y);

            // start removing from the least stake. Iterate until we know enough have been removed.
            let mut removed = 0;
            for (maybe_index, _stake) in voters_sorted.iter().map(|(who, stake)| (nominator_index(&who), stake)) {
                let index = maybe_index.ok_or(OffchainElectionError::NominatorSnapshotCorrupt)?;
                if compact.remove_voter(index) {
                    crate::log!(
                        trace,
                        "💸 removed a voter at index {} with stake {:?} from compact to reduce the size",
                        index,
                        _stake,
                    );
                    removed += 1
                }

                if removed >= to_remove {
                    break;
                }
            }

            crate::log!(
                warn,
                "💸 {} nominators out of {} had to be removed from compact solution due to size limits.",
                removed,
                compact.voter_count() + removed,
            );
            Ok(compact)
        },
        _ => {
            // nada, return as-is
            crate::log!(info, "💸 Compact solution did not get trimmed due to block weight limits.",);
            Ok(compact)
        },
    }
}

/// Takes an election result and spits out some data that can be submitted to the chain.
///
/// This does a lot of stuff; read the inline comments.
pub fn prepare_submission<T: Config>(
    assignments: Vec<Assignment<T::AccountId, OffchainAccuracy>>,
    winners: Vec<(T::AccountId, ExtendedBalance)>,
    do_reduce: bool,
    maximum_weight: Weight,
) -> Result<(Vec<ValidatorIndex>, CompactAssignments, ElectionScore, ElectionSize), OffchainElectionError> {
    // make sure that the snapshot is available.
    let snapshot_validators = <Module<T>>::snapshot_validators().ok_or(OffchainElectionError::SnapshotUnavailable)?;
    let snapshot_nominators = <Module<T>>::snapshot_nominators().ok_or(OffchainElectionError::SnapshotUnavailable)?;

    // all helper closures that we'd ever need.
    let nominator_index = |a: &T::AccountId| -> Option<NominatorIndex> {
        snapshot_nominators
            .iter()
            .position(|x| x == a)
            .and_then(|i| <usize as TryInto<NominatorIndex>>::try_into(i).ok())
    };
    let validator_index = |a: &T::AccountId| -> Option<ValidatorIndex> {
        snapshot_validators
            .iter()
            .position(|x| x == a)
            .and_then(|i| <usize as TryInto<ValidatorIndex>>::try_into(i).ok())
    };
    let nominator_at = |i: NominatorIndex| -> Option<T::AccountId> { snapshot_nominators.get(i as usize).cloned() };

    let validator_at = |i: ValidatorIndex| -> Option<T::AccountId> { snapshot_validators.get(i as usize).cloned() };

    // both conversions are safe; snapshots are not created if they exceed.
    let size = ElectionSize {
        validators: snapshot_validators.len() as ValidatorIndex,
        nominators: snapshot_nominators.len() as NominatorIndex,
    };

    // Clean winners.
    let winners = sp_npos_elections::to_without_backing(winners);

    // convert into absolute value and to obtain the reduced version.
    let mut staked = sp_npos_elections::assignment_ratio_to_staked(assignments, <Module<T>>::slashable_balance_of_fn());

    // reduce
    if do_reduce {
        reduce(&mut staked);
    }

    // Convert back to ratio assignment. This takes less space.
    let low_accuracy_assignment =
        sp_npos_elections::assignment_staked_to_ratio_normalized(staked).map_err(|e| OffchainElectionError::from(e))?;

    // compact encode the assignment.
    let compact = CompactAssignments::from_assignment(low_accuracy_assignment, nominator_index, validator_index)
        .map_err(|e| OffchainElectionError::from(e))?;

    // potentially reduce the size of the compact to fit weight.
    let maximum_allowed_voters =
        maximum_compact_len::<<T as Config>::WeightInfo>(winners.len() as u32, size, maximum_weight);

    crate::log!(
        debug,
        "💸 Maximum weight = {:?} // current weight = {:?} // maximum voters = {:?} // current votes = {:?}",
        maximum_weight,
        <T as Config>::WeightInfo::submit_solution_better(
            size.validators.into(),
            size.nominators.into(),
            compact.voter_count() as u32,
            winners.len() as u32,
        ),
        maximum_allowed_voters,
        compact.voter_count(),
    );

    let compact = trim_to_weight::<T, _>(maximum_allowed_voters, compact, &nominator_index)?;

    // re-compute the score. We re-create what the chain will do. This is a bit verbose and wastes
    // CPU time, but it is necessary to ensure that the score that we claim is the same as the one
    // calculated by the chain.
    let score = {
        let compact = compact.clone();
        let assignments = compact.into_assignment(nominator_at, validator_at).unwrap();
        let staked =
            sp_npos_elections::assignment_ratio_to_staked(assignments.clone(), <Module<T>>::slashable_balance_of_fn());

        let support_map =
            to_support_map::<T::AccountId>(&winners, &staked).map_err(|_| OffchainElectionError::ElectionFailed)?;
        support_map.evaluate()
    };

    // winners to index. Use a simple for loop for a more expressive early exit in case of error.
    let mut winners_indexed: Vec<ValidatorIndex> = Vec::with_capacity(winners.len());
    for w in winners {
        if let Some(idx) = snapshot_validators.iter().position(|v| *v == w) {
            let compact_index: ValidatorIndex = idx.try_into().map_err(|_| OffchainElectionError::InvalidWinner)?;
            winners_indexed.push(compact_index);
        } else {
            return Err(OffchainElectionError::InvalidWinner);
        }
    }

    Ok((winners_indexed, compact, score, size))
}

#[cfg(test)]
mod test {
    #![allow(unused_variables)]
    use super::*;
    use crate::ElectionSize;

    struct Staking;

    impl crate::WeightInfo for Staking {
        fn bond() -> Weight {
            unimplemented!()
        }
        fn bond_extra() -> Weight {
            unimplemented!()
        }
        fn unbond() -> Weight {
            unimplemented!()
        }
        fn withdraw_unbonded_update(s: u32) -> Weight {
            unimplemented!()
        }
        fn withdraw_unbonded_kill(s: u32) -> Weight {
            unimplemented!()
        }
        fn validate() -> Weight {
            unimplemented!()
        }
        fn nominate(n: u32) -> Weight {
            unimplemented!()
        }
        fn chill() -> Weight {
            unimplemented!()
        }
        fn set_payee() -> Weight {
            unimplemented!()
        }
        fn set_controller() -> Weight {
            unimplemented!()
        }
        fn set_validator_count() -> Weight {
            unimplemented!()
        }
        fn force_no_eras() -> Weight {
            unimplemented!()
        }
        fn force_new_era() -> Weight {
            unimplemented!()
        }
        fn force_new_era_always() -> Weight {
            unimplemented!()
        }
        fn set_invulnerables(v: u32) -> Weight {
            unimplemented!()
        }
        fn force_unstake(s: u32) -> Weight {
            unimplemented!()
        }
        fn cancel_deferred_slash(s: u32) -> Weight {
            unimplemented!()
        }
        fn payout_stakers_dead_controller(n: u32) -> Weight {
            unimplemented!()
        }
        fn payout_stakers_alive_staked(n: u32) -> Weight {
            unimplemented!()
        }
        fn rebond(l: u32) -> Weight {
            unimplemented!()
        }
        fn set_history_depth(e: u32) -> Weight {
            unimplemented!()
        }
        fn reap_stash(s: u32) -> Weight {
            unimplemented!()
        }
        fn new_era(v: u32, n: u32) -> Weight {
            unimplemented!()
        }
        fn submit_solution_better(v: u32, n: u32, a: u32, w: u32) -> Weight {
            (0 * v + 0 * n + 1000 * a + 0 * w) as Weight
        }
        fn kick(w: u32) -> Weight {
            unimplemented!()
        }
    }

    #[test]
    fn find_max_voter_binary_search_works() {
        let size = ElectionSize { validators: 0, nominators: 10 };

        assert_eq!(maximum_compact_len::<Staking>(0, size, 0), 0);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 1), 0);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 999), 0);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 1000), 1);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 1001), 1);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 1990), 1);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 1999), 1);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 2000), 2);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 2001), 2);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 2010), 2);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 2990), 2);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 2999), 2);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 3000), 3);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 3333), 3);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 5500), 5);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 7777), 7);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 9999), 9);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 10_000), 10);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 10_999), 10);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 11_000), 10);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 22_000), 10);

        let size = ElectionSize { validators: 0, nominators: 1 };

        assert_eq!(maximum_compact_len::<Staking>(0, size, 0), 0);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 1), 0);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 999), 0);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 1000), 1);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 1001), 1);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 1990), 1);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 1999), 1);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 2000), 1);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 2001), 1);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 2010), 1);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 3333), 1);

        let size = ElectionSize { validators: 0, nominators: 2 };

        assert_eq!(maximum_compact_len::<Staking>(0, size, 0), 0);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 1), 0);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 999), 0);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 1000), 1);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 1001), 1);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 1999), 1);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 2000), 2);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 2001), 2);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 2010), 2);
        assert_eq!(maximum_compact_len::<Staking>(0, size, 3333), 2);
    }
}
